# -*- encoding: utf-8 -*-
"""
========================================
@Time   ：2021/9/6 11:11
@Auther ：shutao
@File   ：QNN.py
@IDE    ：PyCharm
@Github ：https://github.com/NameLacker
@Gitee  ：https://gitee.com/nameLacker
========================================
"""

import numpy as np
import paddle
from paddle_quantum.circuit import UAnsatz
from paddle_quantum.utils import pauli_str_to_matrix

# 超参数设置
num_qubits = 2
theta_size = 4
ITR = 60  # 设置迭代次数
LR = 0.4  # 设置学习率
SEED = 999  # 固定随机种子
paddle.seed(SEED)


class vqe_demo(paddle.nn.Layer):
    def __init__(self, shape, dtype='float64'):
        super(vqe_demo, self).__init__()

        # 初始化一个长度为theta_size的可学习参数列表，并用[0, 2*pi]的均匀分布来填充初始值
        self.theta = self.create_parameter(shape=shape,
                                           default_initializer=paddle.nn.initializer.Uniform(low=0., high=2 * np.pi),
                                           dtype=dtype, is_bias=False)

    def forward(self, H_info):
        # 初始量子电路
        cir = UAnsatz(num_qubits)

        # 添加量子门
        cir.ry(self.theta[0], 0)
        cir.ry(self.theta[1], 1)
        cir.cnot([0, 1])
        cir.ry(self.theta[2], 0)
        cir.ry(self.theta[3], 1)

        # 选择用量子态的向量表示
        cir.run_state_vector()

        # 计算当前量子态下关于观测值H_info的期望值
        # 也就是 <psi|H|psi>
        loss = cir.expecval(H_info)
        return loss


def run_train():
    # 首先生成泡利字符串表示下的哈密顿量
    # 相当于0.4*kron(I, Z) + 0.4*kron(Z, I) + 0.2*kron(X, X)
    # 其中， X，Y, Z是泡利矩阵， I是单位矩阵
    H_info = [[0.4, 'z0'], [0.4, 'z1'], [0.2, 'x0,x1']]

    # 把记录的关于哈密顿量的信息转化为矩阵表示
    H_matrix = pauli_str_to_matrix(H_info, num_qubits)

    # 定义网络
    vqe = vqe_demo([theta_size])

    # 优化器
    opt = paddle.optimizer.Adam(learning_rate=LR, parameters=vqe.parameters())

    loss_list = []
    parameter_list = []
    # 优化循环
    for itr in range(ITR):
        # 前向传播
        loss = vqe(H_info)

        # 反向传播
        loss.backward()
        opt.minimize(loss)
        opt.clear_grad()

        # 记录学习曲线
        loss_list.append(loss.numpy()[0])
        parameter_list.append(vqe.parameters()[0].numpy())
        if itr % 10 == 0:
            print('iter:', itr, '  loss: %.4f' % loss.numpy())

    print('计算得到的基态能量是: ', loss_list[-1])
    print('真实的基态能量为: ', np.linalg.eigh(H_matrix)[0][0])


if __name__ == '__main__':
    run_train()
